import cv2 as cv

def add_logo_to_video_stream(logo_path, window_name="ZHL", window_size=(1080, 960), camera_index=0):
    """
    在实时视频流中添加一个 logo 图像。

    参数:
        logo_path (str): logo 图像的路径。
        window_name (str): OpenCV 窗口的名称。
        window_size (tuple): 窗口的大小，格式为 (宽度, 高度)。
        camera_index (int): 摄像头的索引号，默认为 0（第一个摄像头）。
    """
    # 加载 logo 图像
    img_logo = cv.imread(logo_path)
    if img_logo is None:
        print(f"Error: Logo image not found at {logo_path}!")
        return

    # 获取实时视频流
    cap = cv.VideoCapture(camera_index)
    if not cap.isOpened():
        print("Error: Camera not found!")
        return

    # 设置窗口大小
    cv.namedWindow(window_name)
    cv.resizeWindow(window_name, *window_size)

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Error: Failed to read frame from camera!")
            break

        # 获取 logo 的大小
        h, w, _ = img_logo.shape

        # 检查 logo 大小是否适合视频帧
        frame_height, frame_width, _ = frame.shape
        if w > frame_width or h > frame_height:
            print("Error: Logo size exceeds frame size!")
            break

        # 从背景中切割出和 logo 一样的大小的区域
        img_bg_roi = frame[:h, :w]

        # logo 进行灰度处理，并二值化
        img_logo_gray = cv.cvtColor(img_logo, cv.COLOR_BGR2GRAY)
        # 创建白色 logo 的掩膜，目的是获取 logo 的颜色，有 logo 没背景（反阈值法）
        _, white = cv.threshold(img_logo_gray, 170, 255, cv.THRESH_BINARY_INV)
        fg1 = cv.bitwise_and(img_logo, img_logo, mask=white)
        # 创建黑色 logo 的掩膜，目的是获取到背景颜色，有背景没 logo（阈值法）
        _, black = cv.threshold(img_logo_gray, 170, 255, cv.THRESH_BINARY)
        fg2 = cv.bitwise_and(img_bg_roi, img_bg_roi, mask=black)

        # 图像融合
        img_bg_roi[:] = fg1 + fg2

        # 显示整个视频帧
        cv.imshow(window_name, frame)

        if cv.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放资源
    cap.release()
    cv.destroyAllWindows()

# 调用函数
if __name__ == '__main__':
    add_logo_to_video_stream('txycl/images/baolong.png', window_name="ZHL", window_size=(1080, 960), camera_index=0)